#include <windows.h>
#include <stdio.h>
#include "api.h"
#include "image.h"
#include "png.h"
#include "zlib.h"

#pragma comment (lib, "zdll.lib")
#pragma pack(1)

typedef struct {
	FILE *fpng;
	PBYTE data;
	int line;
	int dimx, dimy;
	PBYTE ptr;
} *PPNGWRITERHANDLER, PNGWRITERHANDLER;

BYTE WRITE_PNG_SIGNATURE[] = {137, 80, 78, 71, 13, 10, 26, 10, 0};

/* Table of CRCs of all 8-bit messages. */
unsigned long crc_table[256];

/* Flag: has the table been computed? Initially false. */
int crc_table_computed = 0;

/* Make the table for a fast CRC. */
void pngWriteMakeCrcTable(void)
{
 unsigned long c;
 int n, k;

 for (n = 0; n < 256; n++) {
   c = (unsigned long) n;
   for (k = 0; k < 8; k++) {
     if (c & 1)
       c = 0xedb88320L ^ (c >> 1);
     else
       c = c >> 1;
   }
   crc_table[n] = c;
 }
 crc_table_computed = 1;
}

/* Update a running CRC with the bytes buf[0..len-1]--the CRC
  should be initialized to all 1's, and the transmitted value
  is the 1's complement of the final running CRC (see the
  crc() routine below)). */

unsigned long pngWriteUpdateCrc(unsigned long crc, unsigned char *buf,
                        int len)
{
	unsigned long c = crc;
	int n;

	if (!crc_table_computed){
		pngWriteMakeCrcTable();
	}
	for (n = 0; n < len; n++) {
		c = crc_table[(c ^ buf[n]) & 0xff] ^ (c >> 8);
	}
	return c;
}

unsigned long pngWriteReverseNumber(unsigned long number){
	unsigned long value = 0;
	for(int i=0;i<4;i++){
		value <<= 8;
		value += number & 0x000000FF;
		number >>= 8;
	}
	return value;
}

void pngWriteChunk(FILE *fpng, long id, long len, void *data){
	crc_table_computed = 0;
	unsigned long crc = pngWriteUpdateCrc(0xffffffffL, (BYTE*)&id, 4);
	if (len){
		crc = pngWriteUpdateCrc(crc, (BYTE*)data, len);
	}
	crc = crc ^ 0xffffffffL;
	crc = pngWriteReverseNumber(crc);
	long realLen = pngWriteReverseNumber(len);
	fwrite(&realLen, 1, 4, fpng);
	fwrite(&id, 1, 4, fpng);
	if (len){
		fwrite(data, 1, len, fpng);
	}
	fwrite(&crc, 1, 4, fpng);
}

void pngWriteHeaderChunk(PPNGWRITERHANDLER data){
	PNGHEADER header;
	header.width = pngWriteReverseNumber(data->dimx);
	header.height = pngWriteReverseNumber(data->dimy);
	header.compression = 0;
	header.filter = 0;
	header.colorType = 2;
	header.bitdepth = 8;
	header.interlaced = 0;
	pngWriteChunk(data->fpng, ID_IDHR, sizeof(PNGHEADER), &header);
}

PPNGWRITERHANDLER pngWriteCreateHandler(PIMAGEHEADER header){
	NEW_VALUE(PNGWRITERHANDLER, handler);
	handler->dimx = header->dx;
	handler->dimy = header->dy;
	handler->line = header->dx*3+1;
	handler->data = (PBYTE)malloc(handler->line*header->dy);
	handler->ptr = handler->data;
	return handler;
}

void pngWriteDeleteHandler(PPNGWRITERHANDLER handler){
	free(handler->data);
	free(handler);
}

void pngWriteCloseImage(PPNGWRITERHANDLER handler){
	fclose(handler->fpng);
}

void pngWriteLine(PPNGWRITERHANDLER handler, PBYTE input){
	PBYTE output = handler->ptr;
	*output = 0;	// No filter the lines in the target PNG
	output++;
	for(int i=0;i<handler->dimx;i++){
		output[0] = input[2];
		output[1] = input[1];
		output[2] = input[0];
		input+=3;
		output+=3;
	}
	handler->ptr += handler->line;
}

int pngWriteDeflate(BYTE *in, int lenIn, BYTE *out, int lenOut){
	z_stream strm;
	int ret;
	int len;
	memset(&strm, 0, sizeof(strm));
	ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
	strm.avail_in = lenIn;
	strm.next_in = in;
	strm.avail_out = lenOut;
	strm.next_out = out;
	deflate(&strm, Z_FINISH);
	len = strm.total_out;
	deflateEnd(&strm);
	return len;
}

void pngWriteDataBlocks(PPNGWRITERHANDLER data, PBYTE ptr, int total){
	int len = (total<1024) ? total : (total >> 10);
	int current = total;
	while(current>len){
		pngWriteChunk(data->fpng, ID_IDAT, len, ptr);
		current -= len;
		ptr += len;
	}
	if (current){
		pngWriteChunk(data->fpng, ID_IDAT, current, ptr);
	}
}

void pngWriteDataChunk(PPNGWRITERHANDLER data){
	int total = data->dimy*data->line;
	PBYTE raw = (PBYTE)malloc(total);
	int lenRaw = pngWriteDeflate(data->data, total, raw, total);
	pngWriteDataBlocks(data, raw, lenRaw);
	free(raw);
}

DWORD pngWrite(UINT command, PVOID handler, PVOID param){
	PPNGWRITERHANDLER data = (PPNGWRITERHANDLER)handler;
	switch(command){
	case IMAGE_SAVER_CREATEHANDLER:
		return (DWORD)pngWriteCreateHandler((PIMAGEHEADER)param);
	case IMAGE_SAVER_CREATEIMAGE:
		data->fpng = fopen((char*)param, "wb+");
		fwrite(WRITE_PNG_SIGNATURE, 8, 1, data->fpng);
		break;
	case IMAGE_SAVER_SAVEHEADER:
		pngWriteHeaderChunk(data);
		break;
	case IMAGE_SAVER_WRITELINE:
		pngWriteLine(data, (PBYTE)param);
		break;
	case IMAGE_SAVER_CLOSEIMAGE:
		pngWriteDataChunk(data);
		pngWriteChunk(data->fpng, ID_IEND, 0, NULL);
		pngWriteCloseImage(data);
		break;
	case IMAGE_SAVER_DELETEHANDLER:		
		pngWriteDeleteHandler(data);
		break;
	}
	return 0;
}